{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import datetime as dt\n",
    "import copy\n",
    "import tushare as ts\n",
    "import math\n",
    "import os\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "XIAOTIAN YANG\n",
    "Quantitative Finance Internship Works\n",
    "#2018\n",
    "This file is merely for calculating the historical index weights for CSI300/CSI500/CSI50 approximately\n",
    "It can' t fullfill the requirements for accurate uses.\n",
    "In order to use this .py file, you have to have your data files prepared already with exact data format as required in the following."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All codes in the next cell are about how to read data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CNE5S_100_Asset_Exposure = pd.read_table(read_path_input+'CNE5S_100_Asset_Exposure.'+target_date_rt, sep = '|', header=2)\n",
    "# CNE5S_100_Asset_Exposure.rename(columns={'!Barrid':'Barrid'}, inplace = True)\n",
    "\n",
    "# CNE5S_100_DlyFacRet = pd.read_table(read_path_input+'CNE5S_100_DlyFacRet.'+target_date_rt, sep = '|', header=2)\n",
    "# CNE5S_100_DlyFacRet.rename(columns={'!Factor':'Factor'}, inplace = True)\n",
    "\n",
    "# CNE5_Daily_Asset_Price = pd.read_table(read_path_input+'CNE5_Daily_Asset_Price.'+target_date_rt, sep = '|', header=1)\n",
    "# CNE5_Daily_Asset_Price.rename(columns={'!Barrid':'Barrid', 'DlyReturn%':'DlyReturn'}, inplace = True)\n",
    "\n",
    "\n",
    "# CNE5_100_Asset_DlySpecRet = pd.read_table(read_path_input+'CNE5_100_Asset_DlySpecRet.'+target_date_rt, sep = '|', header=2)\n",
    "# CNE5_100_Asset_DlySpecRet.rename(columns={'!Barrid':'Barrid'}, inplace = True)\n",
    "\n",
    "# CHN_LOCALID_Asset_ID = pd.read_csv(read_path_input+'CHN_LOCALID_Asset_ID.'+target_date_rt, sep = '|', header=1)\n",
    "# CHN_LOCALID_Asset_ID.rename(columns={'!Barrid':'Barrid'}, inplace = True)\n",
    "# CHN_LOCALID_Asset_ID = CHN_LOCALID_Asset_ID.dropna()\n",
    "# CHN_LOCALID_Asset_ID = lastestBarrid(CHN_LOCALID_Asset_ID)\n",
    "\n",
    "# CNE5S_100_Asset_Exposure = pd.read_table(read_path_input+'CNE5S_100_Asset_Exposure.'+target_date_rt, sep = '|', header=2)\n",
    "# CNE5S_100_Asset_Exposure.rename(columns={'!Barrid':'Barrid'}, inplace = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class PassingError(Exception): \n",
    "        \"\"\"Error encountered while passing arguments.\"\"\"\n",
    "        pass\n",
    "    \n",
    "\n",
    "def getwigt(alist):\n",
    "    epylist = len(alist)*[None]\n",
    "    for ii in range(len(alist)):\n",
    "        epylist[ii] = alist[ii][1]\n",
    "    return epylist\n",
    "\n",
    "\n",
    "def getwigt_name(alist):\n",
    "    epylist = len(alist)*[None]\n",
    "    for ii in range(len(alist)):\n",
    "        epylist[ii] = alist[ii][0]\n",
    "    return epylist\n",
    "\n",
    "def getwigt_localid(alist):\n",
    "    epylist = len(alist)*[None]\n",
    "    for ii in range(len(alist)):\n",
    "        epylist[ii] = alist[ii][2]\n",
    "    return epylist\n",
    "\n",
    "def latest_id(local_id,iddf):\n",
    "    dfa = iddf[iddf.AssetID.apply(lambda x: x == 'CN' + local_id[:6])]\n",
    "    result = dfa[dfa.EndDate == max(dfa.EndDate)].values[0,0]\n",
    "    return result\n",
    "\n",
    "def Del_Narows(tgt_df, col_name):\n",
    "    nalist = list(tgt_df.index[pd.isnull(tgt_df[col_name])].values)\n",
    "    for na in nalist:\n",
    "        tgt_df.drop(na,inplace = True)\n",
    "    return tgt_df\n",
    "\n",
    "def make_dict(keylist,vllist):\n",
    "    emptydict = dict()\n",
    "    if len(keylist) == len(vllist):\n",
    "        for i in range(len(keylist)):\n",
    "            emptydict[keylist[i]] = vllist[i]\n",
    "        return emptydict\n",
    "    else:\n",
    "        print(\"Two lists are not the same long\")\n",
    "        \n",
    "def lcl_bra(df1, df2):\n",
    "    lcldict = []\n",
    "    nanlist = []\n",
    "    for index, row in df1.iterrows():\n",
    "        try:\n",
    "            barraid = latest_id(row[1], df2)\n",
    "        except:\n",
    "            print('\\n' + '\\033[1;31m' + row[1] + ' This is not included in barra' + '\\033[0m')\n",
    "            nanlist.append(row[0])\n",
    "        else:\n",
    "            lcldict.append([barraid,row[-1],row[1]])\n",
    "    return lcldict, nanlist\n",
    "\n",
    "def lcl_bra_init(df1, df2):\n",
    "    lcldict = []\n",
    "    nanlist = []\n",
    "    for index, row in df1.iterrows():\n",
    "        try:\n",
    "            barraid = latest_id(row[0], df2)\n",
    "        except:\n",
    "            print('\\n' + '\\033[1;31m' + row[1] + ' This is not included in barra' + '\\033[0m')\n",
    "            nanlist.append(row[0])\n",
    "        else:\n",
    "            lcldict.append([barraid,row[-1],row[0]])\n",
    "    return lcldict, nanlist\n",
    "\n",
    "def Get_dfbylist_test(targetdf, key_nm, keylist, target_nm):\n",
    "    emptylist = [0]*len(keylist)\n",
    "    \n",
    "    for i in range(len(keylist)):\n",
    "        rowlist = [0,0]\n",
    "        rowlist[0] = keylist[i]\n",
    "        rowlist[1] = ((targetdf[targetdf[key_nm] == keylist[i]])[target_nm]).values[0]\n",
    "        emptylist[i] = rowlist\n",
    "        \n",
    "    return emptylist\n",
    "\n",
    "def Get_dfbylist(targetdf, key_nm, keylist, target_nm):\n",
    "    emptylist = [0]*len(keylist)\n",
    "    \n",
    "    for i in range(len(keylist)):\n",
    "        rowlist = [0]*2\n",
    "        rowlist[0] = keylist[i]\n",
    "        rowlist[1] = ((targetdf[targetdf[key_nm] == keylist[i]])[target_nm]).values[0]\n",
    "        emptylist[i] = rowlist\n",
    "    return emptylist\n",
    "\n",
    "#This is a function use market value to calculate the change of index weights\n",
    "\n",
    "def add_prep(df):\n",
    "    preps = list(df.close)\n",
    "    preps.insert(0,list(df.close)[0])\n",
    "    df['pre_close'] = preps[:-1]\n",
    "    return df\n",
    "\n",
    "def date_cvt(astr):\n",
    "    if len(astr) == 10:\n",
    "        return astr[0:4]+astr[5:7]+astr[8:10]\n",
    "    \n",
    "    elif len(astr) == 8:\n",
    "        return astr[0:4]+'-'+astr[4:6]+'-'+astr[6:8]\n",
    "    \n",
    "    else:\n",
    "        raise PassingError('The passed argument is not as expected')\n",
    "\n",
    "def Wigt_Cal_Mrt(local_idpair, wigts_crt_list, crt_index, tgt_index, crt_mrtvle_table, tgt_mrtvle_table, cal_direction):\n",
    "    lgth_wigts = len(wigts_crt_list)\n",
    "    empty_list = [None]*lgth_wigts\n",
    "    temp_list = [None]*lgth_wigts\n",
    "    \n",
    "    if cal_direction == 'back':\n",
    "        for index in range(lgth_wigts):\n",
    "            temp_row = [None]*2\n",
    "            wigt_pair = wigts_crt_list[index]\n",
    "            idx_change = crt_index/tgt_index          \n",
    "            crt_mrtvle = crt_mrtvle_table[crt_mrtvle_table['Barrid'] == wigt_pair[0]].Capt\n",
    "            tgt_mrtvle = tgt_mrtvle_table[tgt_mrtvle_table['Barrid'] == latest_id(wigt_pair[2],local_idpair)].Capt\n",
    "            idvl_mrtvle_change = (crt_mrtvle.values[0]/\n",
    "                                      tgt_mrtvle.values[0])\n",
    "\n",
    "            temp_row[0] = wigt_pair[2]\n",
    "            temp_row[1] = (wigt_pair[1])*(idx_change)/(idvl_mrtvle_change)\n",
    "            temp_list[index] = temp_row\n",
    "        empty_list = [[x[0],(x[1]*100/sum(getwigt(temp_list)))] for x in temp_list]\n",
    "            \n",
    "    elif cal_direction == 'forward':\n",
    "        \n",
    "        for index in range(lgth_wigts):\n",
    "            temp_row = [None]*2\n",
    "            wigt_pair = wigts_crt_list[index]\n",
    "\n",
    "            idx_change = tgt_index/crt_index\n",
    "            crt_mrtvle = crt_mrtvle_table[crt_mrtvle_table['Barrid'] == wigt_pair[0]].Capt\n",
    "            tgt_mrtvle = tgt_mrtvle_table[tgt_mrtvle_table['Barrid'] == latest_id(wigt_pair[2],local_idpair)].Capt\n",
    "            idvl_mrtvle_change = (tgt_mrtvle.values[0]/\n",
    "                                      crt_mrtvle.values[0])\n",
    "            temp_row[0] = wigt_pair[2]\n",
    "            temp_row[1] = wigt_pair[1]/(idx_change)*(idvl_mrtvle_change)\n",
    "            temp_list[index] = temp_row\n",
    "        empty_list = [[x[0],(x[1]*100/sum(getwigt(temp_list)))] for x in temp_list]\n",
    "\n",
    "    else:\n",
    "        raise PassingError('The passed argument is not as expected')\n",
    "        \n",
    "    return empty_list\n",
    "\n",
    "\n",
    "def Weight_CalwithDate(start_triweight, start_date, trdys_list, index_table_input, calc_dicn, read_path_input):\n",
    "    dly_rtdict = dict()\n",
    "    for date in trdys_list:\n",
    "        daily_return_table = Del_Narows(pd.read_table(read_path_input +'CNE5_Daily_Asset_Price.' +\\\n",
    "                                                              date, sep='|', header=1),'Capt')\n",
    "        daily_return_table.rename(columns={'!Barrid': 'Barrid', 'DlyReturn%':'DlyReturn'}, inplace=True)\n",
    "        dly_rtdict[date] = daily_return_table\n",
    "    \n",
    "    wigt_dict = dict()\n",
    "    \n",
    "    start_dtfmt2 = date_cvt(start_date)\n",
    "    start_mrkt = dly_rtdict[start_date]\n",
    "\n",
    "    for datea in (trdys_list):\n",
    "        date_fmt2 = date_cvt(datea)\n",
    "        \n",
    "        localid_asset_id = pd.read_csv(read_path_input + 'CHN_LOCALID_Asset_ID.' + datea, sep='|', header=1)\n",
    "        localid_asset_id.rename(columns={'!Barrid': 'Barrid'}, inplace=True)\n",
    "\n",
    "        tgt_index_fc = index_table_input[index_table_input.date == date_fmt2].close.values[0]\n",
    "\n",
    "        crt_index_fc = index_table_input[index_table_input.date == start_dtfmt2].close.values[0]\n",
    "        \n",
    "        temp_result = Wigt_Cal_Mrt(localid_asset_id, start_triweight, crt_index_fc, tgt_index_fc, start_mrkt, dly_rtdict[datea], calc_dicn)\n",
    "\n",
    "        wigt_dict[datea] = temp_result\n",
    "        \n",
    "        print(datea+' calculated')\n",
    "        \n",
    "    return (wigt_dict,dly_rtdict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mainfunc(start_date, end_date, read_path_input,wight_file_path,index_type):\n",
    "    \n",
    "    start_date_dt = dt.datetime.strptime(start_date, '%Y%m%d')\n",
    "    end_date_dt = dt.datetime.strptime(end_date, '%Y%m%d')\n",
    "\n",
    "    if (start_date_dt - end_date_dt).days < 0:\n",
    "        calc_dicn = 'forward'\n",
    "        min_date = start_date\n",
    "        max_date = end_date\n",
    "\n",
    "    elif (start_date_dt - end_date_dt).days > 0:\n",
    "        calc_dicn = 'back'\n",
    "        max_date = start_date\n",
    "        min_date = end_date\n",
    "\n",
    "    else:\n",
    "        raise PassingError('Start Date and End Date must be different')\n",
    "\n",
    "    trade_days_index = ts.get_k_data(code = '000300', index = True,\\\n",
    "                                  start = date_cvt(min_date), end = date_cvt(max_date))\n",
    "\n",
    "    trade_dates_list = list(trade_days_index.date.apply(lambda x : date_cvt(x)))\n",
    "    print(trade_dates_list)\n",
    "    chn_localid_asset_id = pd.read_csv(read_path_input + 'CHN_LOCALID_Asset_ID.' + start_date, sep='|', header=1)\n",
    "    chn_localid_asset_id.rename(columns={'!Barrid': 'Barrid'}, inplace=True)\n",
    "#read the initial weight file\n",
    "\n",
    "    if index_type == 'hs300':\n",
    "        idxPrc = add_prep(ts.get_k_data(code = '000300', index = True, start = '2014-12-01')).copy()[['date','close','pre_close']]\n",
    "        indexWigt = pd.read_csv(wight_file_path + 'hs300_' + start_date+'.csv')\n",
    "    elif index_type == 'zz500':\n",
    "        idxPrc = add_prep(ts.get_k_data(code = '000905', index = True, start = '2014-12-01')).copy()[['date','close','pre_close']]\n",
    "        indexWigt = pd.read_csv(wight_file_path + 'zz500_' + start_date+'.csv')\n",
    "    elif index_type == 'sz50':\n",
    "        idxPrc = add_prep(ts.get_k_data(code = '000016', index = True, start = '2014-12-01')).copy()[['date','close','pre_close']]\n",
    "        indexWigt = pd.read_csv(wight_file_path + 'sz50_' + start_date+'.csv')\n",
    "    else:\n",
    "        raise PassingError('The passed argument is not as expected')\n",
    "        \n",
    "    indexWeight = indexWigt[['localid','weight']]\n",
    "#     indexWeight = indexWigttp.rename(columns={'Date': 'date', 'Constituent_Code':'code','Constituent_Name':'name','Weight':'weight'})\n",
    "    nmwigt_pair,nanlist = lcl_bra_init(Del_Narows(indexWeight,'weight'),chn_localid_asset_id)\n",
    "\n",
    "    indexPrice = Del_Narows(idxPrc,'close')\n",
    "    # indexPrice.drop(len(idxPrc)-1,inplace=True)\n",
    "    indexPrice['close'] = idxPrc['close'].astype('float64')\n",
    "    indexPrice['pre_close'] = idxPrc['pre_close'].astype('float64')\n",
    "\n",
    "    index300_daily_weight,dly_rt = Weight_CalwithDate(nmwigt_pair, start_date, trade_dates_list,\\\n",
    "                                                  indexPrice, calc_dicn, read_path_input)\n",
    "    return (index300_daily_weight,dly_rt)\n",
    "\n",
    "\n",
    "def func_csv(index_daily_weight, pathin, indxmode):\n",
    "    if indxmode == 'hs300':\n",
    "        for exc in index_daily_weight:\n",
    "            wgtdf = pd.DataFrame(data = np.array([getwigt_name(index_daily_weight[exc]),\\\n",
    "                                              getwigt(index_daily_weight[exc])]).T, columns=['localid','weight'])\n",
    "            wgtdf.to_csv(pathin+'hs300w_'+exc+'.csv',index = None)\n",
    "\n",
    "    if indxmode == 'zz500':\n",
    "        for exc in index_daily_weight:\n",
    "            wgtdf = pd.DataFrame(data = np.array([getwigt_name(index_daily_weight[exc]),\\\n",
    "                                              getwigt(index_daily_weight[exc])]).T, columns=['localid','weight'])\n",
    "            wgtdf.to_csv(pathin+'zz500w_'+exc+'.csv',index = None)\n",
    "            \n",
    "    if indxmode == 'sz50':\n",
    "        for exc in index_daily_weight:\n",
    "            wgtdf = pd.DataFrame(data = np.array([getwigt_name(index_daily_weight[exc]),\\\n",
    "                                              getwigt(index_daily_weight[exc])]).T, columns=['localid','weight'])\n",
    "            wgtdf.to_csv(pathin+'sz50w_'+exc+'.csv',index = None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a big dictionary containing all related dates in your calculation\n",
    "date_dict = dict()\n",
    "all_dates = ts.get_k_data(code = '000300', index = True,\\\n",
    "                            start = date_cvt('20171201'), end = date_cvt('20180631')).date\n",
    "monthlist = []\n",
    "\n",
    "# list all months and set up innitial varaibles\n",
    "mmlist = ['1712','1801','1802','1803','1804','1805','1806']\n",
    "headlist = ['hs300','sz50','zz500']\n",
    "pathin = 'D:/datasets/'\n",
    "pathout = 'out_weight/'\n",
    "wigt_file_path = 'weight_in/'\n",
    "\n",
    "# make integers for date_dict as keys and lists of dates within a certain month as values\n",
    "for month in range(len(mmlist)):\n",
    "    date_dict[month] = [date_cvt(x) for x in list(all_dates[all_dates.apply(lambda x: date_cvt(x)[:6] == '20'+mmlist[month])])]\n",
    "\n",
    "for m in range(1,len(mmlist)):\n",
    "    for h in headlist:\n",
    "        startdt = date_dict[m-1][-1]\n",
    "        enddt = date_dict[m][-2]\n",
    "        indxmode = h\n",
    "        print(startdt, enddt, pathin, mmlist[m], indxmode)\n",
    "        index_daily_weight, dly_rt = mainfunc(startdt, enddt, pathin, wigt_file_path, indxmode)\n",
    "        func_csv(index_daily_weight, pathout, indxmode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
